package com.zszq.ms.bigmodeltest.model;

import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 工厂类
 *
 * @author 孟帅
 * @since 2025/3/31
 */
@Component
public class BaseModelFactory implements ApplicationContextAware {
    /**
     * 模型map
     */
    private final Map<String, BaseModel> modelMap = new ConcurrentHashMap<>();


    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        modelMap.clear();
        applicationContext.getBeansOfType(BaseModel.class).forEach((k, v) -> {
            modelMap.put(v.getModelCode(), v);
        });
    }

    /**
     * 获取模型
     *
     * @param modelCode 模型编码
     * @return 模型
     */
    public BaseModel getModel(String modelCode) throws Exception {
        BaseModel baseModel = modelMap.get(modelCode);
        if (baseModel == null) {
            throw new Exception("模型不存在");
        }
        return baseModel;
    }
}
